#include "pch.h"
#include "LogManager.h"
#include "..\Common\SRWLock.h"
#include "..\Common\Utility.h"

wchar_t CLogManager::m_szLogPath[];

CLogManager::CLogManager()
	: m_bWriting(false)
	, m_bFirstRun(true)
	, m_pTpTimer(NULL)
{
	AssertInUIThread();

	DWORD dwSize = GetTempPathW(ARRAYSIZE(m_szLogPath), m_szLogPath);
	if (dwSize == 0 || dwSize > MAX_PATH - 40)
		*m_szLogPath = 0;

	PathAppendW(m_szLogPath, L"GaojieDrive");
}

CLogManager::~CLogManager()
{
	AssertInUIThread();
	assert(m_pTpTimer == NULL);
	assert(m_mapLogFiles.empty());
}

bool CLogManager::Init()
{
	AssertInUIThread();

	m_pTpTimer = CreateThreadpoolTimer([](PTP_CALLBACK_INSTANCE, PVOID pContext, PTP_TIMER)
	{
		((CLogManager *)pContext)->OnTimer();
	}, this, NULL);

	if (m_pTpTimer == NULL)
		return false;

	// Write every 200ms
	FILETIME ftDueTime;
	(INT64&)ftDueTime = -200 * 10000;
	SetThreadpoolTimer(m_pTpTimer, &ftDueTime, 200, 0);

	return true;
}

void CLogManager::Exit()
{
	AssertInUIThread();

	if (m_pTpTimer == NULL)
		return;

	// Stop and wait for timer to exit
	SetThreadpoolTimer(m_pTpTimer, NULL, 0, 0);
	WaitForThreadpoolTimerCallbacks(m_pTpTimer, TRUE);
	CloseThreadpoolTimer(m_pTpTimer);
	m_pTpTimer = NULL;

	CExclusiveLock lock(m_SRWLock);
	auto mapLogFiles = std::move(m_mapLogFiles);
	lock.Unlock();	// Unlock before destructing mapLogFiles
}

void CLogManager::WriteLine(LPCVOID lpAddress, const std::string& strLine)
{
	{
		// Access map using shared lock
		CSharedLock lock(m_SRWLock);

		auto it = m_mapLogFiles.find(lpAddress);
		if (it != m_mapLogFiles.end())
		{
			it->second.WriteLine(strLine);
			return;
		}
	}

	// Modify map using exclusive lock
	CExclusiveLock lock(m_SRWLock);

	auto it = m_mapLogFiles.find(lpAddress);
	if (it != m_mapLogFiles.end())
	{
		it->second.WriteLine(strLine);
		return;
	}

	// Get module handle from address
	HMODULE hModule;
	if (!GetModuleHandleExW(GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS | GET_MODULE_HANDLE_EX_FLAG_UNCHANGED_REFCOUNT, (LPCWSTR)lpAddress, &hModule))
	{
		assert(false);
		return;
	}

	// Add a new CLogFile
	DWORD dwImageSize = PIMAGE_NT_HEADERS((LPCSTR)hModule + PIMAGE_DOS_HEADER(hModule)->e_lfanew)->OptionalHeader.SizeOfImage;
	it = m_mapLogFiles.emplace(std::piecewise_construct, std::forward_as_tuple(hModule, dwImageSize), std::forward_as_tuple(hModule)).first;
	it->second.WriteLine(strLine);
}

void CLogManager::OnTimer()
{
	AssertInThreadpool();

	if (m_bFirstRun.exchange(false))
		DelExpiredLogs();

	if (m_bWriting.exchange(true))
		return;

	// Get CLogFile and unlock before FlushCache
	std::vector<CLogFile *> vecLogFiles;
	CSharedLock lock(m_SRWLock);

	for (auto& [key, logFile] : m_mapLogFiles)
		vecLogFiles.emplace_back(&logFile);

	lock.Unlock();

	for (const auto& pLogFile : vecLogFiles)
		pLogFile->FlushCache();

	m_bWriting = false;
}

// Delete log files older than 7 days
void CLogManager::DelExpiredLogs()
{
	AssertInThreadpool();

	wchar_t szFileName[MAX_PATH];
	PathCombineW(szFileName, m_szLogPath, L"*.log");

	FILETIME ftDueTime;
	GetSystemTimeAsFileTime(&ftDueTime);
	(INT64&)ftDueTime -= 7 * 86400 * 10000000i64;	// 7 days ago

	Util::EnumDirectory(szFileName, [&szFileName, ftDueTime](FIND_DATA&& findData)
	{
		if ((INT64&)findData.ftLastWriteTime < (INT64&)ftDueTime)
		{
			PathRemoveFileSpecW(szFileName);
			PathAppendW(szFileName, findData.cFileName);

			if (findData.dwFileAttributes & FILE_ATTRIBUTE_READONLY)
				SetFileAttributesW(szFileName, FILE_ATTRIBUTE_NORMAL);

			DeleteFileW(szFileName);
		}

		return true;
	});
}
